В предыдущем параграфе рассматривались равномерные оценки разницы истинного и эмпирического рисков. Если в рассматриваемом классе моделей есть «плохие», то равномерные оценки становятся слишком пессимистичными. Часто нельзя гарантировать, что что наш алгоритм обучения их никогда не выбирает, поэтому класс моделей $\mathcal{F}$ для равномерной оценки не получится сузить до класса только «хороших» моделей. Но можно надеяться, что плохие выучиваются не слишком часто. Например, известно, что градиентный спуск обычно сходится к хорошим моделям (об этом мы ещё поговорим в параграфе про implicit bias). В этом параграфе мы разберём элегантный способ учесть «предпочтения» алгоритма обучения в оценке разницы рисков.

Вспомним равномерную оценку для конечного $\mathcal{F}$:

$$\mathbb{P}\left(\sup_{f\in\mathcal{F}} (R(f) - \hat R_m(f)) \geq \epsilon\right) = \mathbb{P}\left(\exists f\in\mathcal{F}: \; (R(f) - \hat R_m(f)) \geq \epsilon\right) \leq$$

$$\leq \sum_{f\in\mathcal{F}} \mathbb{P}(R(f) - \hat R_m(f) \geq \epsilon) \leq \vert\mathcal{F}\vert e^{-2 m \epsilon^2} \quad \forall \epsilon > 0,$$

где $\vert\mathcal{F}\vert$ – мощность класса $\mathcal{F}$. Эта оценка формально верна и для бесконечного $\mathcal{F}$, но смысл её теряется. Давайте попробуем исправить это.

Пусть $\mathcal{F}$ не более, чем счётно. Для каждого $f \in \mathcal{F}$ возьмём своё $\epsilon(f)$. Если взять $\epsilon(f)$ таким, чтобы $\sum_{f \in \mathcal{F}} e^{-2 m \epsilon^2(f)}$ было конечным, то приходим к осмысленной оценке:

$$\mathbb{P}\left(\exists f\in\mathcal{F}: \; (R(f) - \hat R_m(f)) \geq \epsilon(f)\right) \leq$$

$$\sum_{f\in\mathcal{F}} \mathbb{P}\left(R(f) - \hat R_m(f) \geq \epsilon(f)\right) \leq \sum_{f \in \mathcal{F}} e^{-2 m \epsilon^2(f)}.$$

Рассмотрим теперь некоторое вероятностное распределение $P(f)$ на $\mathcal{F}$. В качестве $\epsilon(f)$ возьмём

$$e^{-2 m \epsilon^2(f)} = P(f) e^{-2 m \tilde\epsilon^2},$$

где $\tilde\epsilon \in \mathbb{R}_+$. Из этого уравнения получаем следующее выражение для $\epsilon(f)$:

$$\epsilon(f) = \sqrt{\tilde\epsilon^2 + \frac{1}{2m} \log\frac{1}{P(f)}}.$$

В итоге, для любого $\tilde\epsilon > 0$ получаем оценку:

$$\mathbb{P}\left(\exists f\in\mathcal{F}: \; (R(f) - \hat R_m(f)) \geq \sqrt{\tilde\epsilon^2 + \frac{1}{2m} \log\frac{1}{P(f)}}\right) \leq e^{-2 m \tilde\epsilon^2}.$$

Или, что то же самое, с вероятностью $\geq 1 - \delta$ по $S_m$ для любого $f \in \mathcal{F}$:

$$R(f)-\hat{R}_m(f)\leq\sqrt{\frac{1}{2m}\left(\log\frac{1}{\delta}+\log\frac{1}{P(f)}\right)}$$

Заметим, что если $\mathcal{F}$ конечно, а $P(f)$ – равномерное распределение, то оценка выше совпадает с равномерной оценкой. Если же наш алгоритм обучения предпочитает выбирать модели, для которых $P(f)$ велико, то оценка улучшается по сравнению с равномерной. Таким образом, распределение $P(f)$ «кодирует» наши представления о предпочтениях алгоритма. Будем называть $P(f)$ «априорным распределением».

Как обобщить оценку выше на несчётные классы моделей? В первую очередь, предположим, что наш алгоритм обучения $\mathcal{A}$ стохастичен, а значит, на выходе даёт не одну модель, а распределение:

$$\hat f_m \sim \hat Q_m = \mathcal{A}(S_m).$$

Будем называть это распределение «апостериорным». Такое рассуждение осмысленно, например, для стохастического градиентного спуска: очевидно, что результат его работы на невыпуклой функции потерь недетерминирован (он может сходиться в разные локальные минимумы).

Заметим, что главное отличие апостериорного распределения от априорного в том, что первое зависит от данных, а второе – нет. Важно понимать при этом, что, несмотря на названия, эти два распределения не связаны между собой никаким вариантом формулой Байеса. Сходство с байесовским подходом скорее внешнее. Поэтому слова «априорное» и «апостериорное» имеет смысл писать в кавычках, но для экономии места мы их будем в дальнейшем опускать.

Оценки разности рисков, о которых речь пойдёт ниже, называются PAC-байесовскими (PAC-bayesian, где PAC – probably approximately correct).

Сформулируем одну из классических оценок из этого класса:

Теорема Макаллестера. Пусть $\mathcal{F}$ – множество моделей и $P$ – распределение на $\mathcal{F}$. Тогда для любого $\delta \in (0,1)$ с вероятностью $\geq 1 - \delta$ по $S_m$ имеем:

$$R(\hat Q_m) \leq \hat R_m(\hat Q_m) + \sqrt{\frac{1}{2m-1} \left(\log \frac{4m}{\delta} + \mathrm{KL}(\hat Q_m\;\vert\vert P) \right)},$$

где $R(Q) = \mathbb{E}_{f \sim Q} R(f)$ и $\hat R_m(Q) = \mathbb{E}_{f \sim Q} \hat R_m(f)$.

Видим, что оценка тем лучше, чем ближе апостериорное распределение к априорному. Здесь работает следующая интуиция. Если для большинства обучающих наборов данных апостериорное распределение близко к априорному, то оно почти не зависит от данных, а значит, истинный риск и риск на обучающей выборке должны быть близки с высокой вероятностью. Если же апостериорное зависит от данных сильно, то, скорее всего, модель сильно переобучается, а значит, оценка не может быть хорошей; в нашем случае она велика из-за большой KL-дивергенции.

Для доказательства теоремы нам понадобятся две леммы:

Лемма 1. Для любого распределения $P$ на $\mathcal{F}$ и для любого $\delta \in (0,1)$ с вероятностью $\geq 1 - \delta$ по $S_m$ имеем:

$$\mathbb{E}_{f \sim P} e^{(2m-1) (\Delta_m(f))^2} \leq \frac{4m}{\delta},$$

где $\Delta_m(f) = \vert R(f) - \hat R_m(f)\vert$.

Лемма 2 (лемма Донскера-Вередана, Donsker-Varadhan). Пусть $P$ и $Q$ – вероятностные распределения на множестве $X$. Тогда для любого $h: \; X \to \mathbb{R}$

$$\mathbb{E}_{x\sim Q} h(x) \leq \log\mathbb{E}_{x \sim P} e^{h(x)} + \mathrm{KL}(Q\vert\vert P).$$

Доказательство теоремы

Применим лемму 2 к $X = \mathcal{F}$, $h = (2m-1) \Delta_m^2$ и $Q = \hat Q_m$:

$$\mathbb{E}_{f \sim \hat Q_m} (2m-1) (\Delta_m(f))^2 \leq \log\mathbb{E}_{f \sim P} e^{(2m-1) (\Delta_m(f))^2} + \mathrm{KL}(\hat Q_m\;\vert\vert\;P).$$

Тогда по лемме 1 с вероятностью $\geq 1 - \delta$ по $S_m$ имеем:

$$\mathbb{E}_{f \sim \hat Q_m} (2m-1) (\Delta_m(f))^2 \leq \log\frac{4m}{\delta} + \mathrm{KL}(\hat Q_m\;\vert\vert\;P).$$

Доказательство, таким образом, легко завершается:

$$R(\hat Q_m) - \hat R_m(\hat Q_m) \leq \vert\mathbb{E}_{f\sim \hat Q_m} (R(f) - \hat R_m(f))\vert \leq$$

$$\leq\mathbb{E}_{f\sim \hat Q_m} |R(f) - \hat R_m(f)| = \mathbb{E}_{f\sim \hat Q_m} \Delta_m(f) \leq$$

$$\leq\sqrt{\mathbb{E}_{f\sim \hat Q_m} (\Delta_m(f))^2} \leq$$

$$\sqrt{\frac{1}{2m-1} \left(\log \frac{4m}{\delta} + \mathrm{KL}(\hat Q_m\;\vert\vert\;P) \right)}.$$

Доказательство леммы 2

Мы рассмотрим лишь простой случай, когда и у $P$, и у $Q$ есть плотности, и они нигде не обращаются в ноль.

$$\mathbb{E}_{x \sim Q} h(x) - \mathrm{KL}(Q\vert\vert P) = \mathbb{E}_{x \sim Q} \left(h(x) - \log\left(\frac{q(x)}{p(x)}\right)\right) =$$

$$=\mathbb{E}_{x \sim Q} \log\left(e^{h(x)} \frac{p(x)}{q(x)}\right) \leq \log \mathbb{E}_{x \sim Q} \left(e^{h(x)} \frac{p(x)}{q(x)}\right) = \log \mathbb{E}_{x \sim P} e^{h(x)}.$$

Доказательство леммы 1

Нам понадобится

Неравенство Маркова. Пусть $X$ – неотрицательная случайная величина. Тогда для любого $a > 0$ имеем

$$\mathbb{P}(X \geq a) \leq \frac{\mathbb{E} X}{a}.$$

В качестве $X$ и $a$ из неравенства Маркова возьмём

$$X = \mathbb{E}_{f \sim P} e^{(2m-1) (\Delta_m(f))^2},\qquad a = 4m / \delta$$

Тогда

$$\mathbb{P}\left(\vphantom{\frac14}\mathbb{E}_{f \sim P} e^{(2m-1) (\Delta_m(f))^2} \leq \frac{4m}{\delta}\right) \leq$$

$$\leq\frac{\delta}{4m}\mathbb{E}_{S_m} \mathbb{E}_{f \sim P} e^{(2m-1) (\Delta_m(f))^2}$$

Значит, нам достаточно доказать, что

$$\mathbb{E}_{S_m} \mathbb{E}_{f \sim P} e^{(2m-1) (\Delta_m(f))^2} \leq 4m.$$

Мы докажем даже более сильное соотношение:

$$\mathbb{E}_{S_m}e^{(2m-1)(\Delta_m(f))^2}\leq4m\quad\forall f\in\mathcal{F}$$

Заметим, что из неравенства Хёффдинга будет следовать

$$\mathbb{P}_{S_m}(\Delta_m(f)\geq\epsilon)\leq2e^{-2m\epsilon^2}\quad\forall\epsilon>0\quad\forall f\in\mathcal{F}$$

Для простоты предположим, что распределение $\Delta_m(f)$ имеет плотность для любого $f \in \mathcal{F}$; обозначим её $p_f(\Delta)$. В этом случае мы можем ограничить матожидания по $S_m$ напрямую:

$$\mathbb{E}_{S_m} e^{(2m-1) (\Delta_m(f))^2} = \int_0^\infty e^{(2m-1) \epsilon^2} p_f(\epsilon) \, d\epsilon =$$

$$\int_0^\infty e^{(2m-1) \epsilon^2} \frac{d}{d\epsilon} \left(-\int_\epsilon^\infty p_f(\Delta) \, d\Delta\right) \, d\epsilon =$$

$$= \left.-\left(e^{(2m-1) \epsilon^2} \int_\epsilon^\infty p_f(\Delta) \, d\Delta\right) \right|_{\epsilon=0}^\infty + 2 (2m-1) \int_0^\infty \epsilon\,e^{(2m-1) \epsilon^2} \int_\epsilon^\infty p_f(\Delta) \, d\Delta \, d\epsilon \leq$$

$$\leq \int_0^\infty p_f(\Delta) \, d\Delta + 2 (2m-1) \int_0^\infty \epsilon\,e^{(2m-1) \epsilon^2} \int_\epsilon^\infty p_f(\Delta) \, d\Delta \, d\epsilon \leq$$

$$\leq 2 + 4 (2m-1) \int_0^\infty \epsilon\,e^{(2m-1) \epsilon^2} e^{-2m \epsilon^2} \, d\epsilon =$$

$$2 + 4 (2m-1) \int_0^\infty \epsilon\,e^{-\epsilon^2} \, d\epsilon = 2 + 2 (2m-1) = 4m.$$

В общем случае, мы не можем предполагать наличие плотности у $\Delta_m(f)$. Доказательство в этом случае можно найти в оригинальной работе D. A. McAllester Some pac-bayesian theorems, а также в конспекте лекций автора этого параграфа.

Теорема Макаллестера – не единственная из возможных пак-байесовских оценок. Например, несколько улучшенную версию той же оценки можно найти в работе Bounds for averaging classifiers. Другие оценки подобного типа можно найти в монографии PAC-Bayesian supervised classification: the thermodynamics of statistical learning.

Применение пак-байесовских оценок к детерминированным алгоритмам обучения

Выше были рассмотрены две PAC-байесовские оценки: одна для не более, чем счётного множества моделей, другая – для произвольного. За возможность использования несчётных классов моделей мы заплатили тем, что алгоритм обучения должен быть недетерминированным (для детерминированных алгоритмов KL-дивергенция в Теореме Макаллестера может вырождаться в бесконечность; например, это так, если априорное распределение гауссово). Чаще всего класс моделей $\mathcal{F}$ всё-таки несчетён: например, если это класс всех сетей фиксированной архитектуры, то он индексируется весами, которых несчётное множество. При этом, хотя используемый алгоритм обучения и в самом деле недетерминирован (стохастический градиентный спуск зависит от случайного выбора батчей и от инициализации весов) и теорема Макаллестера выполняется, финальное распределение моделей очень сложно охарактеризовать, и из-за этого непонятно, как считать KL-дивергенцию.

Предположим, что алгоритм обучения всё-таки детерминирован; этого можно добиться, зафиксировав сид генератора случайных чисел при обучении. Как получить осмысленную PAC-байесовскую оценку для детерминированного алгоритма на несчётном множестве моделей?

Мы рассмотрим два способа.

Первый способ – добавить известный шум в финальную модель, выданную детерминированным алгоритмом. Так, для нейронных сетей, результатом работы алгоритма обучения является набор весов. Если добавить в этот набор гауссовский шум, а также в качестве априорного распределения взять гауссовское, то KL-дивергенцию в теореме Макаллестера можно будет посчитать аналитически.

Дисперсию шума в апостериорном распределении тоже можно обучить с помощью градиентного спуска одновременно с весами, тем самым минимизируя правую часть оценки из вышеупомянутой теоремы. Если в найденную модель удастся добавить шум так, чтобы KL-дивергенция значительно уменьшилась, но при этом риск на обучающей выборке не сильно вырос, то оценка на истинный риск получится хорошей.

Это рассуждение связывает PAC-байесовские оценки и гипотезу о том, что «плоские» («широкие») минимумы хорошо обобщают. В самом деле, если минимум «плоский», то в модель из него можно добавить много шума, не испортив качество на обучении. Оценки, основанные на этом принципе, можно найти в работах Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data и A PAC-Bayesian Approach to Spectrally-Normalized Margin Bounds for Neural Networks.

Второй способ состоит в том, чтобы взять дискретное кодирование $c$ и применить дискретную PAC-байесовскую оценку к закодированной модели вместо оригинальной. Обозначим закодированную модель $f$ через $f_c$. Следуя работе Non-vacuous Generalization Bounds at the ImageNet Scale: a PAC-Bayesian Compression Approach, возьмём априорное распределение с массой, убывающей с ростом длины кода:

$$P_c(f_c) = \frac{1}{Z} m(\vert f_c\vert ) 2^{-|f_c|}.$$

Здесь $\vert f_c\vert$ – длина кода модели $f$, $m(k)$ – некоторое вероятностное распределение на $\mathbb{N}$, а $Z$ – нормализующая константа. Тогда KL-дивергенция примет следующий вид:

$$\mathrm{KL}(\delta_{f_c}\vert\vert {P_c}) = \log Z + \vert f_c\vert \log 2 - \log(m(\vert f_c\vert )).$$

Для того, чтобы KL-дивергенция выше была как можно меньше, необходимо, чтобы наш алгоритм обучения на реалистичных данных сходился в модели с маленькой длиной кода. Для этого будем применять наше кодирование не к оригинальной модели, а к сжатой с помощью некоторого алгоритма сжатия. Здесь мы предполагаем, что модели, к которым сходится наш алгоритм обучения, можно сжать с малыми потерями до моделей с малой длиной кода. Другими словами, мы опираемся на предположение, что обученные модели в некоторым смысле «простые».

Если модель параметризована весами $\theta$, типичный алгоритм сжатия выдаст набор $(S,Q,C)$, где

• $S = s_{1:k} \subset [\dim\theta]$ – позиции ненулевых весов;
• $C = c_{1:r} \subset \mathbb{R}$ – «словарь» весов;
• $Q = q_{1:k}$, $q_i \in [r]$ $\forall i \in [k]$ – квантизованные значения весов.

Выход алгоритма будет выглядеть как $\mathcal{C}(\theta)_i = c_{q_j}$, если $i = s_j$, иначе $0$.

Тогда наивное 32-битное кодирование даст следующую длину:

$$\vert\mathcal{C}(\theta)\vert_c = \vert S\vert_c + \vert Q\vert_c + \vert C\vert_c \leq k (\log\dim\theta + \log r) + 32 r.$$

В работе Non-vacuous Generalization Bounds at the ImageNet Scale: a PAC-Bayesian Compression Approach описанный выше способ применяется к модели MobileNet (свёрточной сети, сконструированной специально для мобильных устройств), обученной на наборе данных ImageNet, и получают верхнюю оценку на истинный риск, равную $96.5\%$ (риск случайного угадывания – $99.9\%$). Хотя такой результат и выглядит очень скромным, но это первая осмысленная оценка обобщающей способности реально используемой нейронной сети на реалистичном наборе данных.